#!/usr/bin/env python3
"""
子数据集节点：用于批量配置LoRA训练的角色/风格子集参数
适配DreamBooth/LoRA等多角色训练场景
"""

from typing import List
from .lora_trainer_utils.constants import CHARACTER_LORA_PARAMS_CATEGORY

class CharacterSubsetNode:
    """角色/风格子数据集节点"""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image_dirs": ("STRING_LIST", {"tooltip": "每个角色/风格的图片文件夹"}),
                "class_tokens": ("STRING_LIST", {"tooltip": "每个角色/风格的触发词/标识符"}),
                "num_repeats": ("INT_LIST", {"tooltip": "每个角色/风格的重复次数"}),
                "is_regs": ("BOOL_LIST", {"tooltip": "每个角色/风格是否为正则化图像"}),
            },
            "optional": {
                "caption_extension": ("STRING", {"default": ".caption", "tooltip": "caption扩展名（如需用.txt可修改）"}),
            }
        }

    RETURN_TYPES = ("CHARACTER_SUBSETS",)
    RETURN_NAMES = ("角色子集参数",)
    FUNCTION = "generate_subsets"
    CATEGORY = CHARACTER_LORA_PARAMS_CATEGORY

    def generate_subsets(self, image_dirs: List[str], class_tokens: List[str], num_repeats: List[int], is_regs: List[bool], caption_extension: str = ".caption"):
        if not (len(image_dirs) == len(class_tokens) == len(num_repeats) == len(is_regs)):
            raise ValueError("所有输入列表长度必须一致！")
        subsets = []
        for img_dir, token, repeat, is_reg in zip(image_dirs, class_tokens, num_repeats, is_regs):
            subset = {
                "image_dir": img_dir,
                "class_tokens": token,
                "num_repeats": repeat,
                "is_reg": is_reg,
                "caption_extension": caption_extension
            }
            subsets.append(subset)
        return (subsets,) 


NODE_CLASS_MAPPINGS = {
    "CharacterSubsetNode": CharacterSubsetNode,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "CharacterSubsetNode": "子数据集参数",
} 